'''
https://leetcode.cn/problems/kth-smallest-element-in-a-sorted-matrix/description/
'''
from typing import List


class Solution:
    def kthSmallest(self, matrix: List[List[int]], k: int) -> int:
        n = len(matrix)
        temp = []
        for row in matrix:
            temp.extend(row)
        temp.sort()
        return temp[k-1]

    # 二分
    def kthSmallest2(self, matrix: List[List[int]], k: int) -> int:
        n = len(matrix)

        def f(num):
            i, j = n-1, -1
            res = 0
            while i >= 0:
                while j+1 < n and matrix[i][j+1] <= num:
                    j += 1
                res += j + 1
                i -= 1
            return res

        l, r =  matrix[0][0], matrix[-1][-1]
        res = -1
        while l <= r:
            mid = l + r >> 1
            if f(mid) >= k:
                res = mid
                r = mid - 1
            else:
                l = mid + 1
        return res

matrix = [[1, 5, 9], [10, 11, 13], [12, 13, 15]]
k = 8
print(*matrix, sep='\n')
print(Solution().kthSmallest2(matrix, k))
